import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
# from ..registry import BACKBONES
from einops import rearrange
from torch.autograd import Variable
from .layers import NTUDGraph

class CTRGC(nn.Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 rel_reduction=8,
                 mid_reduction=1):
        super(CTRGC, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        if in_channels == 3 or in_channels == 9:
            self.rel_channels = 8
            self.mid_channels = 16
        else:
            self.rel_channels = in_channels // rel_reduction
            self.mid_channels = in_channels // mid_reduction
        self.conv1 = nn.Conv2d(self.in_channels,
                               self.rel_channels,
                               kernel_size=1)
        self.conv2 = nn.Conv2d(self.in_channels,
                               self.rel_channels,
                               kernel_size=1)
        self.conv3 = nn.Conv2d(self.in_channels,
                               self.out_channels,
                               kernel_size=1)
        self.conv4 = nn.Conv2d(self.rel_channels,
                               self.out_channels,
                               kernel_size=1)
        self.tanh = nn.Tanh()

    def forward(self, x, A=None, alpha=1):
        x1, x2, x3 = self.conv1(x).mean(-2), self.conv2(x).mean(-2), self.conv3(x)
        x1 = self.tanh(x1.unsqueeze(-1) - x2.unsqueeze(-2))
        x1 = self.conv4(x1) * alpha + (
            A.unsqueeze(0).unsqueeze(0) if A is not None else 0)  # N,C,V,V
        x1 = torch.einsum('ncuv,nctv->nctu', x1, x3)
        return x1


class TemporalConv(nn.Module):

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 stride=1,
                 dilation=1):
        super(TemporalConv, self).__init__()
        pad = (kernel_size + (kernel_size - 1) * (dilation - 1) - 1) // 2
        self.conv = nn.Conv2d(in_channels,
                              out_channels,
                              kernel_size=(kernel_size, 1),
                              padding=(pad, 0),
                              stride=(stride, 1),
                              dilation=(dilation, 1))

        self.bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        return x


class MultiScale_TemporalConv(nn.Module):

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=3,
                 stride=1,
                 dilations=[1, 2, 3, 4],
                 residual=True,
                 residual_kernel_size=1):

        super(MultiScale_TemporalConv, self).__init__()
        assert out_channels % (len(dilations) + 2) == 0, \
            '# out channels should be multiples of # branches'

        # Multiple branches of temporal convolution
        self.num_branches = len(dilations) + 2
        branch_channels = out_channels // self.num_branches
        if type(kernel_size) == list:
            assert len(kernel_size) == len(dilations)
        else:
            kernel_size = [kernel_size] * len(dilations)
        # Temporal Convolution branches
        self.branches = nn.ModuleList([
            nn.Sequential(
                nn.Conv2d(in_channels,
                          branch_channels,
                          kernel_size=1,
                          padding=0),
                nn.BatchNorm2d(branch_channels),
                nn.ReLU(),
                TemporalConv(branch_channels,
                             branch_channels,
                             kernel_size=ks,
                             stride=stride,
                             dilation=dilation),
            ) for ks, dilation in zip(kernel_size, dilations)
        ])

        # Additional Max & 1x1 branch
        self.branches.append(
            nn.Sequential(
                nn.Conv2d(in_channels, branch_channels, kernel_size=1, padding=0),
                nn.BatchNorm2d(branch_channels),
                nn.ReLU(),
                nn.MaxPool2d(kernel_size=(3, 1), stride=(stride, 1), padding=(1, 0)),
                nn.BatchNorm2d(branch_channels))
        )

        self.branches.append(
            nn.Sequential(
                nn.Conv2d(in_channels, branch_channels, kernel_size=1,
                          padding=0, stride=(stride, 1)),
                nn.BatchNorm2d(branch_channels))
        )

        # Residual connection
        if not residual:
            self.residual = lambda x: 0
        elif (in_channels == out_channels) and (stride == 1):
            self.residual = lambda x: x
        else:
            self.residual = TemporalConv(in_channels,
                                         out_channels,
                                         kernel_size=residual_kernel_size,
                                         stride=stride)

    def forward(self, x):
        # Input dim: (N,C,T,V)
        res = self.residual(x)
        branch_outs = []
        for tempconv in self.branches:
            out = tempconv(x)
            branch_outs.append(out)

        out = torch.concat(branch_outs, dim=1)
        out += res
        return out


class unit_tcn(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=9, stride=1):
        super(unit_tcn, self).__init__()
        pad = int((kernel_size - 1) / 2)
        self.conv = nn.Conv2d(in_channels,
                              out_channels,
                              kernel_size=(kernel_size, 1),
                              padding=(pad, 0),
                              stride=(stride, 1))
        self.bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        x = self.bn(self.conv(x))
        return x


class unit_gcn(nn.Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 A,
                 coff_embedding=4,
                 adaptive=True,
                 residual=True):
        super(unit_gcn, self).__init__()
        inter_channels = out_channels // coff_embedding
        self.inter_c = inter_channels
        self.out_c = out_channels
        self.in_c = in_channels
        self.adaptive = adaptive
        self.num_subset = A.shape[0]
        self.convs = nn.ModuleList()

        for i in range(self.num_subset):
            self.convs.append(CTRGC(in_channels, out_channels))

        if residual:
            if in_channels != out_channels:
                self.down = nn.Sequential(
                    nn.Conv2d(in_channels, out_channels, 1),
                    nn.BatchNorm2d(out_channels))
            else:
                self.down = lambda x: x
        else:
            self.down = lambda x: 0

        if self.adaptive:
            self.PA = nn.Parameter(torch.from_numpy(A.astype(np.float32)))
        else:
            # self.A = torch.from_numpy(A.astype(np.float32))
            self.A = Variable(torch.from_numpy(A.astype(np.float32)), requires_grad=False)

        self.alpha = nn.Parameter(torch.zeros(1))
        self.bn = nn.BatchNorm2d(out_channels)
        self.soft = nn.Softmax(-2)
        self.relu = nn.ReLU()

    def forward(self, x):
        y = None
        if self.adaptive:
            A = self.PA
        else:
            A = self.A.cuda(x.get_device())
        for i in range(self.num_subset):
            z = self.convs[i](x, A[i], self.alpha)
            y = z + y if y is not None else z
        y = self.bn(y)
        y += self.down(x)
        y = self.relu(y)
        return y


class TCN_GCN_unit(nn.Module):

    def __init__(self,
                 in_channels,
                 out_channels,
                 A,
                 stride=1,
                 residual=True,
                 adaptive=True,
                 kernel_size=5,
                 dilations=[1, 2]):
        super(TCN_GCN_unit, self).__init__()
        self.gcn1 = unit_gcn(in_channels, out_channels, A, adaptive=adaptive)
        self.tcn1 = MultiScale_TemporalConv(out_channels,
                                            out_channels,
                                            kernel_size=kernel_size,
                                            stride=stride,
                                            dilations=dilations,
                                            residual=False)
        self.relu = nn.ReLU()
        if not residual:
            self.residual = lambda x: 0

        elif (in_channels == out_channels) and (stride == 1):
            self.residual = lambda x: x

        else:
            self.residual = unit_tcn(in_channels,
                                     out_channels,
                                     kernel_size=1,
                                     stride=stride)

    def forward(self, x):
        y = self.relu(self.tcn1(self.gcn1(x)) + self.residual(x))
        return y





# @BACKBONES.register()
class CTRGCN(nn.Module):
    """
    CTR-GCN model from:
    `"Channel-wise Topology Refinement Graph Convolution for Skeleton-Based Action Recognition" <https://arxiv.org/abs/2107.12213>`_
    Args:
        num_point: int, numbers of sketeton point.
        num_person: int, numbers of person.
        base_channel: int, model's hidden dim.
        graph: str, sketeton adjacency matrix name.
        graph_args: dict, sketeton adjacency graph class args.
        in_channels: int, channels of vertex coordinate. 2 for (x,y), 3 for (x,y,z). Default 3.
        adaptive: bool, if adjacency matrix can adaptive.
    """

    def __init__(self,
                 num_point=25,
                 num_person=1,
                 base_channel=64,
                 graph='ntu_rgb_d',
                 graph_args=dict(),
                 in_channels=3,
                 adaptive=True):
        super(CTRGCN, self).__init__()

        if graph == 'ntu_rgb_d':
            self.graph = NTUDGraph(**graph_args)
        else:
            raise ValueError()

        A = self.graph.A  # 3,25,25

        self.num_point = num_point
        self.data_bn = nn.BatchNorm1d(num_person * in_channels * num_point)
        self.base_channel = base_channel

        self.ctrgcn = nn.Sequential(
            TCN_GCN_unit(in_channels, self.base_channel, A, residual=False, adaptive=adaptive),
            TCN_GCN_unit(self.base_channel, self.base_channel, A, adaptive=adaptive),
            TCN_GCN_unit(self.base_channel, self.base_channel, A, adaptive=adaptive),
            TCN_GCN_unit(self.base_channel, self.base_channel, A, adaptive=adaptive),
            TCN_GCN_unit(self.base_channel, self.base_channel * 2, A, stride=2, adaptive=adaptive),
            TCN_GCN_unit(self.base_channel * 2, self.base_channel * 2, A, adaptive=adaptive),
            TCN_GCN_unit(self.base_channel * 2, self.base_channel * 2, A, adaptive=adaptive),
            TCN_GCN_unit(self.base_channel * 2, self.base_channel * 4, A, stride=2, adaptive=adaptive),
            TCN_GCN_unit(self.base_channel * 4, self.base_channel * 4, A, adaptive=adaptive),
            TCN_GCN_unit(self.base_channel * 4, self.base_channel * 4, A, adaptive=adaptive)
        )

        self.pool = nn.AdaptiveAvgPool2d(output_size=(1, 1))

    def forward(self, x):
        N, C, T, V, M = x.shape
        x = rearrange(x, 'b c t v m -> b (m v c) t')
        x = self.data_bn(x)
        x = rearrange(x, 'b (m v c) t -> (b m) c t v', m=M, c=C, v=V)  # [N * M, C, T, V]
        x = self.ctrgcn(x)
        x = self.pool(x)  # NM,C,T,V --> NM,C,1,1
        x = rearrange(x, '(b m) c t v -> b m c t v', m=M).mean(axis=1)  # N,C,1,1
        return x

    def init_weights(self):
        #  among self.modules() Show the layer structure of the network
        for m in self.modules():
            #  nn.Conv2d Layer for processing
            if isinstance(m, (nn.Conv2d, nn.Conv1d)):
                #  use  torch.nn.init.xavier_normal  Method for this layer  weight  To initialize
                torch.nn.init.xavier_normal_(m.weight.data)
                #  And judge whether there is bias (bias), If exist , take  bias  Initialize to full  0
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            # BatchNorm2d Layer for processing
            elif isinstance(m, (nn.BatchNorm2d, nn.BatchNorm1d)):
                # Initialize with constants
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            #  Linear Layer for processing ,Initialize with constants
            elif isinstance(m, nn.Linear):
                #  Initialization of normal distribution , Make the value obey the normal distribution  N(mean, std)
                torch.nn.init.normal_(m.weight.data, 0, 0.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)


if __name__ == '__main__':
    B, C, T, V, M = 2, 3, 300, 25, 1
    x = torch.rand(size=[B, C, T, V, M])
    md = CTRGCN()
    y = md(x)
    print(y.shape)
